import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import asyncio
import websockets
import threading
import json
import time  # 导入time模块，用于心跳包逻辑

class WebSocketClientNode(Node):
    def __init__(self):
        super().__init__('websocket_client_node')
        # 1. ROS 2 发布者：发布到 websocket_cmd 话题
        self.publisher = self.create_publisher(String, 'websocket_cmd', 10)
        self.get_logger().info("WebSocket节点启动，发布话题: websocket_cmd")
        
        # 2. WebSocket 服务器地址
        self.ws_uri = "ws://new_human_server.qupuba.com:9101/ws"
        self.board_name = "body"  # 设备名称
        
        # 3. 启动 WebSocket 客户端线程
        self.ws_thread = threading.Thread(target=self._run_websocket, daemon=True)
        self.ws_thread.start()

    def _run_websocket(self):
        """在独立线程中运行 WebSocket 客户端"""
        asyncio.run(self._connect_and_listen())

    async def _connect_and_listen(self):
        """异步连接 WebSocket 并处理消息"""
        last_heartbeat = 0
        while rclpy.ok():  # 保持与 ROS 2 生命周期一致
            try:
                async with websockets.connect(self.ws_uri) as websocket:
                    self.get_logger().info(f"已连接到 WebSocket 服务器: {self.ws_uri}")
                    # 发送注册消息
                    register_msg = json.dumps({"type": "register", "name": self.board_name})
                    await websocket.send(register_msg)
                    self.get_logger().info(f"发送注册消息: {register_msg}")
                    
                    while rclpy.ok():
                        try:
                            # 带超时的接收，避免阻塞
                            message = await asyncio.wait_for(websocket.recv(), timeout=1.0)
                            self.get_logger().info(f"收到 WebSocket 消息: {message}")
                            # 解析并发布指令到 ROS 2 话题
                            self._handle_websocket_message(message)
                        except asyncio.TimeoutError:
                            # 心跳包逻辑
                            if time.time() - last_heartbeat > 15:
                                heartbeat_msg = json.dumps({"type": "heartbeat"})
                                await websocket.send(heartbeat_msg)
                                last_heartbeat = time.time()
                                self.get_logger().info("发送心跳包")
                        except websockets.exceptions.ConnectionClosed:
                            self.get_logger().error("WebSocket 连接关闭，重试...")
                            break
            except Exception as e:
                self.get_logger().error(f"WebSocket 错误: {e}，5秒后重试...")
                await asyncio.sleep(5)

    def _handle_websocket_message(self, message):
        """解析 WebSocket 消息并发布到 ROS 2 话题"""
        try:
            data = json.loads(message)
        except json.JSONDecodeError:
            self.get_logger().error("WebSocket 消息 JSON 解析错误")
            return
        
        msg_type = data.get("type")
        if not msg_type:
            self.get_logger().error("WebSocket 消息缺少 type 字段")
            return
        
        # 处理 servo_control 类型消息
        if msg_type in ("private", "servo_control"):
            content = data.get("content")
            try:
                payload = json.loads(content) if isinstance(content, str) else content
                self.get_logger().info(f"解析出舵机控制指令: {payload}")
                # 发布指令到 ROS 2 话题（舵机节点会订阅并解析）
                msg = String()
                # 确保payload是列表格式，与串口转换后的格式一致
                if not isinstance(payload, list):
                    payload = [payload]
                msg.data = json.dumps(payload)
                self.publisher.publish(msg)
                self.get_logger().info(f"发布到 websocket_cmd 话题: {msg.data}")
            except Exception as e:
                self.get_logger().error(f"解析舵机控制指令失败: {e}")
        elif msg_type == "heartbeat":
            self.get_logger().info("收到心跳响应")
        else:
            self.get_logger().info(f"忽略未知消息类型: {msg_type}")

def main(args=None):
    rclpy.init(args=args)
    node = WebSocketClientNode()
    try:
        rclpy.spin(node)  # ROS 2 主循环
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()